from fastapi import FastAPI
from fastapi_mcp import FastApiMCP

import requests
import json
import base64

key = ""
# Your API app
api_app = FastAPI()
# ... define your API endpoints on api_app ...
@api_app.get("/")
async def root():
    return {"message": "Hello World"}



@api_app.get("/key/{newkey}")
async def auth(newkey: str):
    """
    这是进行fofa搜索之前所必须的验证部分，
    在运行本mcp服务其他功能之前必须调用该函数，
    在运行其他所有功能之前，必须先验证
    key的内容必须通过询问用户获得
    
    Args:
        key (str): 验证所需要的key

    Returns:
        dict: 返回一个字典，包含key
    """
    global key
    key = newkey
    return {"key": key}

@api_app.get("/search/{content}")
async def search(content: str):
    """
    使用fofa搜索相关内容，
    要求语法规范符合fofa搜索的语法
    
    Args:
        content (str): 需要搜索的内容

    Returns:
        dict: 返回一个字典，包含搜索结果
    """
    

    search_string = content
    encoded_bytes = base64.b64encode(search_string.encode("utf-8"))
    encoded_string = encoded_bytes.decode("utf-8")
    url = "https://fofa.info/api/v1/search/all?&key="+key+"&size=100&qbase64=" + encoded_string
    #print(url)

    res = requests.get(url)
    results = json.loads(res.text)
    #print(results)
    return {"results": results}


# A separate app for the MCP server
mcp_app = FastAPI()

# Create MCP server from the API app
mcp = FastApiMCP(
    api_app,
    base_url="http://localhost:8001",  # The URL where the API app will be running
)

# Mount the MCP server to the separate app
mcp.mount(mcp_app)

# Now you can run both apps separately:
# uvicorn main:api_app --host api-host --port 8001
# uvicorn main:mcp_app --host mcp-host --port 8000